#!/usr/bin/python3
import os
import cv2
import numpy as np
import pandas as pd

# face_feature_list_file = r'/rootfs/media/kasim/Data/data/FaceImage/BoxFeature.txt'
# face_feature_data_file = r'/rootfs/media/kasim/Data/data/FaceImage/Feature.dat'
# face_attr_list_file = r'/rootfs/media/kasim/Data/data/FaceImage/face_detect_list.attr.txt'

face_feature_list_file = r'/rootfs/media/kasim/Data1/data/VideoFaceImage/BoxFeature.txt'
face_feature_data_file = r'/rootfs/media/kasim/Data1/data/VideoFaceImage/Feature.dat'
face_attr_list_file = r'/rootfs/media/kasim/Data1/data/VideoFaceImage/extract_image_list_72.attr.txt'

# VERY_GOOD = 0
# GOOD = 1
# BAD = 2
# VERY_BAD = 3
filter_quality = 3
filter_quality_thr = 0.5

feat_head_offset = 8


def main():
    # file_name;bin_pitch;bin_roll;bin_yaw;pitch;roll;yaw;quality;mask
    attr_cvs = pd.read_csv(face_attr_list_file, sep=';')
    file_names = attr_cvs['file_name']
    qualitys = attr_cvs['quality']
    filter_file_names = set()
    file_count = 0
    for file_name, quality in zip(file_names, qualitys):
        quality = list(map(float, quality.split(',')))
        quality = np.array(quality, dtype=np.float32)
        max_idx = np.argmax(quality)
        if (max_idx == filter_quality) and (quality[max_idx] >= filter_quality_thr):
            filter_file_names.add(file_name)
        file_count += 1
        if file_count % 10000 == 0:
            print('Process Filter File Count: {}'.format(file_count))
    print('Process Filter File Count: {}'.format(file_count))

    with open(face_feature_list_file, 'r') as file:
        file_count = 0
        save_file_count = 0
        file_infos = file.readlines()
        total_file_count = len(file_infos)
        info = file_infos[0].strip().split()
        feat_size = int(info[1].strip().split(',')[-1])
        feat_file = open(face_feature_data_file, 'rb')
        out_feat_file = open(os.path.splitext(face_feature_data_file)[0]+'.sel.dat', 'wb')
        out_file = open(os.path.splitext(face_feature_data_file)[0]+'.sel.txt', 'w')
        for line in file_infos:
            file_count += 1
            if file_count % 10000 == 0:
                print('Process File Count: {}/{}, {}'.format(file_count, total_file_count, save_file_count))
            feat_data = feat_file.read(feat_size)
            info = line.strip().split()
            file_name = info[0]
            if file_name in filter_file_names:
                continue
            no_head_feat_data = feat_data[feat_head_offset:]
            out_feat_file.write(no_head_feat_data)
            out_file.write(file_name+'\n')
            save_file_count += 1
        out_file.close()
        out_feat_file.close()
        feat_file.close()

        print('Process File Count: {}/{}, {}'.format(file_count, total_file_count, save_file_count))


if __name__ == '__main__':
    main()
